package com.tianji.aigc.service;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import redis.clients.jedis.JedisPooled;

import java.time.Instant;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Service
@RequiredArgsConstructor
public class RedisVectorStoreService {

    private final VectorStore vectorStore;
    private final JedisPooled jedis;

    private static final String VECTOR_DOC_ZSET_KEY = "vector_docs"; // 全局索引key
    private static final String VECTOR_DOC_HASH_KEY_PREFIX = "vector_doc:"; // 详情hash前缀

    /**
     * 保存多个文档到向量库，同时写入Redis索引结构
     */
    public void saveDocuments(List<String> messages) {
        List<Document> documents = new ArrayList<>();

        for (String message : messages) {
            String id = UUID.randomUUID().toString();
            long timestamp = Instant.now().toEpochMilli();

            // 存hash元数据
            String hashKey = VECTOR_DOC_HASH_KEY_PREFIX + id;
            jedis.hset(hashKey, Map.of(
                    "id", id,
                    "text", message,
                    "timestamp", String.valueOf(timestamp)
            ));

            // 存到排序集合中
            jedis.zadd(VECTOR_DOC_ZSET_KEY, timestamp, id);

            // 生成 Document
            documents.add(Document.builder()
                    .id(id)
                    .text(message)
                    .metadata(Map.of("timestamp", timestamp))
                    .build());
        }

        // 存入向量数据库
        vectorStore.add(documents);
        log.info("成功保存 {} 条向量数据", documents.size());
    }

    /**
     * 真分页查询（按时间倒序）
     */
    public List<Document> listDocuments(int page, int size) {
        int start = (page - 1) * size;
        int end = start + size - 1;

        var ids = jedis.zrevrange(VECTOR_DOC_ZSET_KEY, start, end);
        if (ids.isEmpty()) return List.of();

        return ids.stream()
                .map(id -> {
                    Map<String, String> map = jedis.hgetAll(VECTOR_DOC_HASH_KEY_PREFIX + id);
                    return Document.builder()
                            .id(id)
                            .text(map.get("text"))
                            .metadata(Map.of("timestamp", map.get("timestamp")))
                            .build();
                })
                .collect(Collectors.toList());
    }

    /**
     * 获取总数
     */
    public long countDocuments() {
        return jedis.zcard(VECTOR_DOC_ZSET_KEY);
    }

    /**
     * 删除文档
     */
    public void deleteDocument(String id) {
        jedis.zrem(VECTOR_DOC_ZSET_KEY, id);
        jedis.del(VECTOR_DOC_HASH_KEY_PREFIX + id);
        vectorStore.delete(List.of(id));
        log.info("已删除文档ID: {}", id);
    }
}
